import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import pandas as pd
import os
import threading

# 确保导入tabulate库用于Markdown支持
try:
    import tabulate
except ImportError:
    print("警告: 未找到tabulate库，Markdown导出可能失败。请运行: pip install tabulate")

class ChatDataProcessor:
    def __init__(self, root):
        self.root = root
        self.root.title("聊天数据处理器")
        self.root.geometry("800x600")
        
        # 数据存储
        self.chat_data = None
        self.names = []
        self.file_path = None
        
        # 创建界面
        self.create_widgets()
        
        # 线程锁，用于安全更新GUI
        self.gui_lock = threading.Lock()
    
    def create_widgets(self):
        # 上部分 - 文件选择
        file_frame = tk.Frame(self.root, pady=10)
        file_frame.pack(fill=tk.X)
        
        tk.Label(file_frame, text="Excel文件:").pack(side=tk.LEFT, padx=10)
        self.file_path_var = tk.StringVar()
        tk.Entry(file_frame, textvariable=self.file_path_var, width=50).pack(side=tk.LEFT, padx=5)
        self.browse_button = tk.Button(file_frame, text="浏览", command=self.browse_file)
        self.browse_button.pack(side=tk.LEFT, padx=5)
        self.load_button = tk.Button(file_frame, text="加载数据", command=self.start_load_thread)
        self.load_button.pack(side=tk.LEFT, padx=5)
        
        # 加载进度条
        self.load_progress = ttk.Progressbar(file_frame, orient=tk.HORIZONTAL, length=100, mode='indeterminate')
        # 初始隐藏
        # self.load_progress.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True) 
        
        # 中间部分 - 人名选择
        name_frame = tk.Frame(self.root, pady=10)
        name_frame.pack(fill=tk.X)
        
        tk.Label(name_frame, text="选择人名:").pack(side=tk.LEFT, padx=10)
        self.name_var = tk.StringVar()
        self.name_combo = ttk.Combobox(name_frame, textvariable=self.name_var, width=20, state="readonly")
        self.name_combo.pack(side=tk.LEFT, padx=5)
        
        # 添加手动输入人名的部分
        tk.Label(name_frame, text="或手动输入:").pack(side=tk.LEFT, padx=5)
        self.manual_name_var = tk.StringVar()
        self.manual_name_entry = tk.Entry(name_frame, textvariable=self.manual_name_var, width=15)
        self.manual_name_entry.pack(side=tk.LEFT, padx=5)
        
        self.extract_button = tk.Button(name_frame, text="提取数据", command=self.start_extract_thread)
        self.extract_button.pack(side=tk.LEFT, padx=5)

        # 提取进度条
        self.extract_progress = ttk.Progressbar(name_frame, orient=tk.HORIZONTAL, length=100, mode='indeterminate')
        # 初始隐藏
        # self.extract_progress.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True)
        
        # 下部分 - 数据预览
        preview_frame = tk.Frame(self.root, pady=10)
        preview_frame.pack(fill=tk.BOTH, expand=True)
        
        tk.Label(preview_frame, text="数据预览:").pack(anchor=tk.W, padx=10)
        
        # 创建表格
        self.tree_frame = tk.Frame(preview_frame)
        self.tree_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)
        
        self.tree_scroll_y = tk.Scrollbar(self.tree_frame)
        self.tree_scroll_y.pack(side=tk.RIGHT, fill=tk.Y)
        
        self.tree_scroll_x = tk.Scrollbar(self.tree_frame, orient=tk.HORIZONTAL)
        self.tree_scroll_x.pack(side=tk.BOTTOM, fill=tk.X)
        
        self.tree = ttk.Treeview(self.tree_frame, yscrollcommand=self.tree_scroll_y.set, 
                                xscrollcommand=self.tree_scroll_x.set)
        
        self.tree_scroll_y.config(command=self.tree.yview)
        self.tree_scroll_x.config(command=self.tree.xview)
        
        self.tree.pack(fill=tk.BOTH, expand=True)
        
        # 状态栏
        status_frame = tk.Frame(self.root)
        status_frame.pack(fill=tk.X, side=tk.BOTTOM)
        
        self.status_var = tk.StringVar()
        self.status_var.set("准备就绪")
        tk.Label(status_frame, textvariable=self.status_var, bd=1, relief=tk.SUNKEN, anchor=tk.W).pack(fill=tk.X)

    def start_load_thread(self):
        file_path = self.file_path_var.get()
        if not file_path or not os.path.exists(file_path):
            messagebox.showerror("错误", "请选择有效的Excel文件")
            return
        
        # 禁用按钮，显示进度条
        self.load_button.config(state=tk.DISABLED)
        self.browse_button.config(state=tk.DISABLED)
        self.extract_button.config(state=tk.DISABLED) # 加载时也禁用提取
        self.load_progress.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True) # 显示
        self.load_progress.start()
        self.status_var.set("正在加载数据...")

        # 启动加载线程
        thread = threading.Thread(target=self.load_data, args=(file_path,))
        thread.daemon = True # 主程序退出时线程也退出
        thread.start()

    def load_data(self, file_path):
        try:
            temp_chat_data = pd.read_excel(file_path)
            
            # --- 检查必需的列 --- 
            required_columns = ['消息发送者', '消息接收者']
            missing_columns = [col for col in required_columns if col not in temp_chat_data.columns]
            if missing_columns:
                # 使用after在主线程中显示错误消息
                self.root.after(0, self.show_load_error, f"加载失败：Excel文件中缺少必要的列: {', '.join(missing_columns)}.\n请确保文件包含 '消息发送者' 和 '消息接收者' 列。")
                return # 提前退出线程
            
            # --- 成功加载，准备更新GUI --- 
            sender_names = temp_chat_data['消息发送者'].fillna('').astype(str).unique().tolist()
            receiver_names = temp_chat_data['消息接收者'].fillna('').astype(str).unique().tolist()
            all_names = set(sender_names + receiver_names)
            names = sorted([name for name in all_names if name])
            
            # 使用after在主线程更新GUI
            self.root.after(0, self.update_gui_after_load, temp_chat_data, names)
            
        except Exception as e:
            # 使用after在主线程中显示错误消息
            self.root.after(0, self.show_load_error, f"加载数据失败: {str(e)}")
        finally:
            # 确保无论成功失败都停止进度条并恢复按钮（在主线程中）
            self.root.after(0, self.finalize_load_ui)

    def show_load_error(self, message):
        messagebox.showerror("错误", message)
        self.chat_data = None
        self.name_combo['values'] = []
        self.name_var.set("")
        self.update_preview_table(pd.DataFrame())
        self.status_var.set("加载失败")

    def update_gui_after_load(self, loaded_data, names):
         with self.gui_lock:
            self.chat_data = loaded_data
            self.names = names
            self.status_var.set(f"成功加载数据: {len(self.chat_data)} 行")
            self.name_combo['values'] = self.names
            if self.names:
                self.name_combo.current(0)
                self.extract_button.config(state=tk.NORMAL) # 成功加载后启用提取按钮
            else:
                self.name_var.set("")
                messagebox.showinfo("提示", "未在 '消息发送者' 或 '消息接收者' 列中找到有效的人名。")
                self.extract_button.config(state=tk.DISABLED)
            self.update_preview_table(self.chat_data.head(100))

    def finalize_load_ui(self):
        self.load_progress.stop()
        self.load_progress.pack_forget() # 隐藏
        self.load_button.config(state=tk.NORMAL)
        self.browse_button.config(state=tk.NORMAL)
        # 提取按钮的状态由 update_gui_after_load 决定

    def browse_file(self):
        file_path = filedialog.askopenfilename(
            filetypes=[("Excel文件", "*.xlsx;*.xls")]
        )
        if file_path:
            self.file_path_var.set(file_path)
            self.file_path = file_path

    def update_preview_table(self, data):
        # 清除现有数据
        for item in self.tree.get_children():
            self.tree.delete(item)
        
        # 重置列
        self.tree['columns'] = list(data.columns)
        self.tree.column("#0", width=0, stretch=tk.NO)
        self.tree.heading("#0", text="")
        
        # 设置列
        for col in data.columns:
            self.tree.column(col, anchor=tk.W, width=100)
            self.tree.heading(col, text=col, anchor=tk.W)
        
        # 添加数据行
        for i, row in data.iterrows():
            values = [row[col] for col in data.columns]
            self.tree.insert("", tk.END, text=i, values=values)

    def start_extract_thread(self):
        if self.chat_data is None:
            messagebox.showerror("错误", "请先加载数据")
            return
        
        # 优先使用手动输入的人名，如果为空则使用下拉框选择的人名
        manual_name = self.manual_name_var.get().strip()
        combo_name = self.name_var.get().strip()
        
        if manual_name:
            selected_name = manual_name
        elif combo_name:
            selected_name = combo_name
        else:
            messagebox.showerror("错误", "请选择或输入一个人名")
            return
        
        # 弹出保存对话框，获取保存路径
        save_path = filedialog.asksaveasfilename(
            initialfile=f"{selected_name}聊天记录.md", 
            defaultextension=".md",             
            filetypes=[("Markdown 文件", "*.md"), ("CSV 文件", "*.csv"), ("所有文件", "*.*")]
        )
        
        if not save_path: # 如果用户取消保存
            return

        # 禁用按钮，显示进度条
        self.extract_button.config(state=tk.DISABLED)
        self.load_button.config(state=tk.DISABLED) # 提取时也禁用加载
        self.browse_button.config(state=tk.DISABLED)
        self.extract_progress.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True) # 显示
        self.extract_progress.start()
        self.status_var.set(f"正在为 {selected_name} 提取数据...")

        # 启动提取线程
        thread = threading.Thread(target=self.extract_data, args=(selected_name, save_path))
        thread.daemon = True
        thread.start()

    def extract_data(self, selected_name, save_path):
        try:
            # 筛选逻辑：现在使用 str.contains() 进行部分匹配
            # 只要'消息发送者'或'消息接收者'包含所选名字，就提取该行
            filtered_data = self.chat_data[
                (self.chat_data['消息发送者'].fillna('').astype(str).str.contains(selected_name)) | 
                (self.chat_data['消息接收者'].fillna('').astype(str).str.contains(selected_name))
            ].copy() # 使用copy避免SettingWithCopyWarning
            
            record_count = len(filtered_data)
            
            if record_count == 0:
                # 使用after在主线程显示提示
                self.root.after(0, messagebox.showinfo, "提示", f"没有找到与 {selected_name} 相关的聊天记录")
            else:
                # 更新预览 (在主线程)
                self.root.after(0, self.update_preview_table, filtered_data.head(100))
                
                # 根据文件扩展名决定保存方式
                try:
                    if save_path.lower().endswith('.md'):
                        # 尝试使用to_markdown，依赖tabulate
                        try:
                            markdown_output = filtered_data.to_markdown(index=False)
                            with open(save_path, 'w', encoding='utf-8') as f:
                                f.write(markdown_output)
                        except Exception as md_error:
                            # 如果to_markdown失败，使用备选方案或提供有用的错误消息
                            print(f"Markdown导出失败: {str(md_error)}")
                            
                            # 备选方案：使用自定义实现
                            try:
                                # 简单的自定义Markdown表格生成
                                markdown_content = []
                                
                                # 添加表头
                                header = '| ' + ' | '.join(filtered_data.columns) + ' |'
                                separator = '| ' + ' | '.join(['---' for _ in filtered_data.columns]) + ' |'
                                markdown_content.append(header)
                                markdown_content.append(separator)
                                
                                # 添加数据行
                                for _, row in filtered_data.iterrows():
                                    row_str = '| ' + ' | '.join([str(cell) for cell in row]) + ' |'
                                    markdown_content.append(row_str)
                                
                                with open(save_path, 'w', encoding='utf-8') as f:
                                    f.write('\n'.join(markdown_content))
                            except Exception as custom_md_error:
                                # 如果自定义实现也失败，则改为CSV
                                csv_path = save_path.replace('.md', '.csv')
                                filtered_data.to_csv(csv_path, index=False, encoding='utf-8-sig')
                                error_message = f"Markdown导出失败，已改为CSV保存到 {csv_path}"
                                self.root.after(0, messagebox.showwarning, "警告", error_message)
                                save_path = csv_path
                    else:
                        # 其他格式，默认使用CSV
                        filtered_data.to_csv(save_path, index=False, encoding='utf-8-sig')
                    
                    # 显示成功消息
                    success_message = f"已成功导出 {record_count} 条记录到 {save_path}"
                    self.root.after(0, messagebox.showinfo, "成功", success_message)
                    self.root.after(0, self.status_var.set, f"成功为 {selected_name} 导出 {record_count} 条记录")
                except Exception as e:
                    error_message = f"文件保存失败: {str(e)}"
                    self.root.after(0, messagebox.showerror, "错误", error_message)
                    self.root.after(0, self.status_var.set, "导出失败")

        except Exception as e:
            # 使用after在主线程显示错误消息
            error_message = f"处理数据时出错: {str(e)}"
            self.root.after(0, messagebox.showerror, "错误", error_message)
            self.root.after(0, self.status_var.set, f"为 {selected_name} 提取数据时出错")
        finally:
            # 确保恢复UI (在主线程)
            self.root.after(0, self.finalize_extract_ui)

    def finalize_extract_ui(self):
        self.extract_progress.stop()
        self.extract_progress.pack_forget() # 隐藏
        self.extract_button.config(state=tk.NORMAL)
        # 只有在数据加载成功时才恢复加载按钮
        if self.chat_data is not None:
            self.load_button.config(state=tk.NORMAL)
        self.browse_button.config(state=tk.NORMAL)


if __name__ == "__main__":
    root = tk.Tk()
    app = ChatDataProcessor(root)
    root.mainloop()
